package nl.cypherpunk.statefuzzer.smtp;
import nl.cypherpunk.statefuzzer.tls.TLSClient;
import nl.cypherpunk.statefuzzer.tls.TLSSession;
import nl.cypherpunk.statefuzzer.tls.TLSTestService;

import java.io.*;
import java.net.*;
import java.nio.MappedByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.charset.StandardCharsets;
import java.util.*;

public class SMTPTestService {
    Socket socket;
    OutputStream output;
    InputStream input;

    String host = "127.0.0.1";
    int port = 25;
    String hostname = "localhost";
    // Act as a TLS client
    boolean ROLE_CLIENT = true;
    // Restart server after every session
    boolean REQUIRE_RESTART = true;
    // Timeout in ms
    int RECEIVE_MSG_TIMEOUT = 100;
    // Send output from TLS implementation to console
    boolean CONSOLE_OUTPUT = false;

    String cmd;
    public Process targetProcess;

    String redirectFile;
    ArrayList<byte[]> currentCase;
    int totalExecs;
    int prevCoverage;
    String outputDir;
    String env_cmd="";
    PrintWriter out;
    BitSet coverageMap;
    Double initScore;
    String SendUser="ubuntu@ubuntu";
    String SendPassword="xlj932834897";
    String ReceiveUser="ubuntu@ubuntu";
    MappedByteBuffer fileCoverage;
    HashSet<String> crashHash;
    String version = "";
    public SMTPTestService() throws Exception {

    }

    public void setTarget(String target) throws Exception {
        if(target.equals("server")) {
            ROLE_CLIENT = true;
        }
        else if(target.equals("client")) {
            ROLE_CLIENT = false;
        }
        else {
            throw new Exception("Unknown target");
        }
    }

    public void setHost(String host) {
        this.host = host;
    }

    public void setPort(int port) {
        this.port = port;
    }

    public void setCommand(String cmd) {
        this.cmd = cmd;
    }

    public void setRestartTarget(boolean restart) {
        this.REQUIRE_RESTART = restart;
    }

    public void setReceiveMessagesTimeout(int timeout) {
        RECEIVE_MSG_TIMEOUT = timeout;
    }

    public void setConsoleOutput(boolean enable) {
        CONSOLE_OUTPUT = enable;
    }

    public void setOutputDir(String dir){
        this.outputDir = dir;
    }
    public void setRedirectFile(String output) {
        redirectFile = output;
    }
    public void setVersion(String version) {
        this.version = version;
    }
    public void start() throws Exception {
        currentCase = new ArrayList<>();
        prevCoverage = 0;
        totalExecs = 0;
        crashHash = new HashSet<>();
        if(ROLE_CLIENT) {
            /*String killall = "sudo killall exim";
            ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
            Process tmp = pbkill.start();
            tmp.waitFor();*/
            String killall = "sudo rm /home/exim"+version;
            ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
            Process tmp = pbkill.start();
            tmp.waitFor();
            if(cmd != null && !cmd.equals("")) {
                ProcessBuilder pb = new ProcessBuilder(cmd.split(" "));
                pb.redirectErrorStream(true);
                //pb.redirectOutput(new File(redirectFile));
                targetProcess = pb.start();
                Thread.sleep(10000);
            }
            //connectSocket();
        }
        else {
        }
        fileCoverage = new RandomAccessFile("/home/exim"+version, "r").getChannel()
                .map(FileChannel.MapMode.READ_ONLY, 0, 40008);
    }

    public void reset() throws Exception {
        //System.out.println("RESET");
        if(socket!=null)
        {
            socket.close();
        }
        totalExecs ++;
        if(totalExecs%1000==0)
            System.out.println(totalExecs);
        if(ROLE_CLIENT) {
            if (!currentCase.isEmpty()) {
                int coverage = getCoverage();
                //reader.close();
                if (coverage == 1)
                    seedSave(currentCase, this.outputDir + File.separator + "Seed" + File.separator + "Crash_" + String.format("%06d", totalExecs) + "_" + String.format("%05d", coverage));
                if (coverage > prevCoverage) {
                    System.out.println(coverage);
                    seedSave(currentCase, this.outputDir + File.separator + "Seed" + File.separator + "A_" + String.format("%06d", totalExecs) + "_" + String.format("%05d", coverage));

                    prevCoverage = coverage;
                }
                currentCase.clear();
            }
            if (REQUIRE_RESTART && cmd != null && !cmd.equals("")) {
                targetProcess.destroy();

                String target = this.outputDir.substring(0, this.outputDir.indexOf("2"));
                String killall = "sudo killall " + target;
                if (REQUIRE_RESTART) {
                    ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
                    Process tmp = pbkill.start();
                    tmp.waitFor();
                }
                Thread.sleep(100);

                ProcessBuilder pb = new ProcessBuilder(cmd.split(" "));
                pb.redirectErrorStream(true);

                targetProcess = pb.start();

                Thread.sleep(5000);
                getCoverage();
            }
            Thread.sleep(100);
            String tmp = "";
            try {
                tmp = connectSocket();
            } catch (Exception e)
            {
                Thread.sleep(1000);
                tmp = connectSocket();
            }
            int i = 0;
            while(!tmp.equals("220"))
            {
                targetProcess.destroy();

                String target = this.outputDir.substring(0, this.outputDir.indexOf("2"));
                String killall = "sudo killall " + target;

                ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
                Process kill = pbkill.start();
                kill.waitFor();

                Thread.sleep(100);

                ProcessBuilder pb = new ProcessBuilder(cmd.split(" "));
                pb.redirectErrorStream(true);

                targetProcess = pb.start();

                Thread.sleep(5000);
                tmp = connectSocket();
            }

        } else {

        }
    }

    public String connectSocket() throws UnknownHostException, IOException,Exception {
        socket = new Socket(host, port);
        socket.setTcpNoDelay(true);
        socket.setSoTimeout(RECEIVE_MSG_TIMEOUT);

        output = socket.getOutputStream();
        input = socket.getInputStream();
        return receiveMessages();
    }

    public void listenSocket() throws UnknownHostException, IOException {
        ServerSocket server = new ServerSocket();
        server.bind(new InetSocketAddress(host, port));
        socket = server.accept();
        socket.setTcpNoDelay(true);
        socket.setSoTimeout(RECEIVE_MSG_TIMEOUT);

        output = socket.getOutputStream();
        input = socket.getInputStream();

        server.close();
    }

    public void closeSocket() throws IOException, Exception {
        socket.close();
    }

    public String receiveMessages() throws Exception {
        String out = "";
        String statusCode = null;
        if(socket.isInputShutdown())
            return "ConnectionClosed";
        try {
            byte[] code = new byte[3];
            input.read(code, 0, 3);
            StringBuffer tStringBuf=new StringBuffer ();
            char[] tChars=new char[3];

            for(int i=0;i<3;i++)
                tChars[i]=(char)code[i];

            tStringBuf.append(tChars);
            out=tStringBuf.toString();

            if(tChars[0]< 0x31 || tChars[0]> 0x40)
                out = "ConnectionClosed";
            if(out.equals("221"))
                out = "221ConnectionClosed";
            if(out.equals("500"))
                out = "554";

        }
        catch (SocketTimeoutException e)
        {
            //socket.close();
            return "Empty";
        }
        /*
        byte[] data = new byte[1000];
        while (input.available() > 0) {
            int length = input.read(data);
        }*/


        if(out.contains("ConnectionClosed")) {
            socket.close();
            return out;
        }
        BufferedReader br = new BufferedReader(new InputStreamReader(input, "utf-8"));
        String msg = "";

        while(input.available()>0)
            msg = br.readLine();
        return out;
    }

    void sendMessage(byte[] msg) throws Exception {
        if(!socket.isOutputShutdown())
            output.write(msg);
    }

    public void close() {
        if(targetProcess != null) {
            targetProcess.destroy();
        }
    }

    public String processSymbol(String input) throws Exception {
        String inAction = input;
        byte[] out = null;

        if(!socket.isConnected() || socket.isClosed()) return "ConnectionClosed";

        try {
            if (inAction.equals("helo")) {
                out = buildHelo();
            } else if(inAction.equals("mail")){
                out = buildMailFrom();
            } else if(inAction.equals("rcpt")){
                out = buildRcptTo();
            } else if(inAction.equals("data")){
                out = buildData();
                /*sendMessage(data);
                currentCase.add(data);
                out = buildMessage();*/
            } else if(inAction.equals("message")){
                out = buildMessage();
            } else if(inAction.equals("rset")){
                out = buildRset();
            } else if(inAction.equals("quit")){
                out = buildQuit();
            } else if(inAction.equals("vrfy")){
                out = buildVRFY();
            } else if(inAction.equals("expn")){
                out = buildEXPN();
            } else if(inAction.equals("bdat")){
                out = buildBDAT();
            } else {
                System.out.println("Unknown input symbol (" + inAction + ")...");
                System.exit(0);
            }
            sendMessage(out);
            currentCase.add(out);
            return receiveMessages();
        }
        catch(SocketException e) {
            //String outAction = "ConnectionClosedException";
            String outAction = "ConnectionClosed";

            return outAction;
        }
    }

    public byte[] buildHelo() throws Exception{
        String message = "HELO " + this.hostname +"\r\n";
        return message.getBytes("US-ASCII");
    }

    public byte[] buildMailFrom() throws Exception{
        String message = "MAIL FROM:<"+this.SendUser+">"+"\r\n";
        return message.getBytes("US-ASCII");
    }

    public byte[] buildRcptTo() throws Exception{
        String message = "RCPT TO:<"+ReceiveUser+">"+"\r\n";
        return message.getBytes("US-ASCII");
    }
    public byte[] buildData() throws Exception{
        //String message = "subject:myxulinjie\r\nfrom:"+this.SendUser+"\r\nto:"+this.ReceiveUser+"\r\nContent-Type: text/plain;\r\ntest\r\n.\r\n";
        String message = "DATA\r\n";
        return message.getBytes("US-ASCII");
    }
    public byte[] buildBDAT() throws Exception{
        //String message = "subject:myxulinjie\r\nfrom:"+this.SendUser+"\r\nto:"+this.ReceiveUser+"\r\nContent-Type: text/plain;\r\ntest\r\n.\r\n";
        String message = "BDAT 10\r\n";
        return message.getBytes("US-ASCII");
    }

    public byte[] buildMessage() throws Exception {
        String message = "subject:myxulinjie\r\nfrom:"+this.SendUser+"\r\nto:"+this.ReceiveUser+"\r\nContent-Type: text/plain;\r\ntest\r\n.\r\n";
        return message.getBytes("US-ASCII");
    }
    public byte[] buildRset() throws Exception{
        String message = "RSET\r\n";
        return message.getBytes("US-ASCII");
    }
    public byte[] buildQuit() throws Exception{
        String message = "QUIT\r\n";
        return message.getBytes("US-ASCII");
    }

    public byte[] buildVRFY() throws Exception{
        String message = "VRFY "+this.SendUser+"\r\n";
        return message.getBytes("US-ASCII");
    }

    public byte[] buildEXPN() throws Exception{
        String message = "EXPN "+this.SendUser+"\r\n";
        return message.getBytes("US-ASCII");
    }

    public String sendFuzzingMessage (byte[] message) throws Exception
    {
        try {
            if (socket.isClosed())
                return "ConnectionClosed";
            output.write(message);
            return receiveFuzzingMessages();
        }
        catch(SocketException e) {
            //String outAction = "ConnectionClosedException";
            String outAction = "ConnectionClosed";

            return outAction;
        }
    }

    public String receiveFuzzingMessages() throws Exception {
        try{
            String out = receiveMessages();
            if(out.contains("ConnectionClosed")) {
                socket.close();
            }
            return out;
        }
        catch(SocketException e) {
            //String outAction = "ConnectionClosedException";
            String outAction = "ConnectionClosed";

            return outAction;
        }

    }

    //change by pany
    public void seedSave(List<byte[]> seed, String filename) {
        BufferedOutputStream bos = null;
        FileOutputStream fos = null;
        File file = null;
        try {
            file = new File(filename);
            if (file.exists()) {
                return;
            }
            fos = new FileOutputStream(file);
            bos = new BufferedOutputStream(fos);
            for (int i = 0; i < seed.size(); i++) {
                byte[] tmp = intToByteArray(seed.get(i).length);
                bos.write(tmp, 0, 4);
                bos.write(seed.get(i), 0, seed.get(i).length);
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (bos != null) {
                try {
                    bos.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
            if (fos != null) {
                try {
                    fos.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
        }
    }

    public static byte[] intToByteArray(int i) {
        byte[] result = new byte[4];
        result[3] = (byte)((i >> 24) & 0xFF);
        result[2] = (byte)((i >> 16) & 0xFF);
        result[1] = (byte)((i >> 8) & 0xFF);
        result[0] = (byte)(i & 0xFF);
        return result;
    }

    public int getCoverage() throws IOException{
        int coverage = 0;
        String resultStr = new String();
        if(ROLE_CLIENT) {
            BufferedReader reader = new BufferedReader(new InputStreamReader(targetProcess.getInputStream()));
            StringBuilder builder = new StringBuilder();
            String line = null;
            try {
                Thread.sleep(50);
                while (reader.ready()) {
                    line = reader.readLine();
                    if (line == null)
                        break;
                    if (line.contains("Coverage")) {
                        try {
                            int index = line.indexOf("Coverage");
                            String cover = line.substring(index + 9).trim();
                            coverage = Integer.parseInt(cover);
                        } catch (NumberFormatException e) {
                            e.printStackTrace();
                        }
                    }
                    if (line.contains("AddressSanitizer"))
                    {
                        coverage = 1;
                        System.out.println(resultStr);
                        System.out.println(line);
                        break;
                    }
                    resultStr = resultStr.concat(line);
                }
            }
            catch(Exception e)
            {
                System.out.println(e);
            }
        }else
        {
            //due to the test case will result in restarting
            if(coverageMap == null)
                coverage = 0;
            else
                coverageMap.cardinality();
            BufferedReader reader = new BufferedReader(new InputStreamReader(targetProcess.getInputStream()));
            StringBuilder builder = new StringBuilder();
            String line = null;
            try{
                while(true) {
                    //Thread.sleep(200);
                    while (reader.ready()) {
                        line = reader.readLine();
                        if (line == null)
                            break;
                        if (line.contains("Bitmap")) {
                            try {
                                int index = line.indexOf("Bitmap");
                                String cover = line.substring(index + 8).trim();
                                if (coverageMap == null) {
                                    coverageMap = new BitSet(cover.length());
                                }
                                for (int i = 0; i < cover.length(); i++)
                                    if (cover.charAt(i) == '1')
                                        coverageMap.set(i);
                                coverage = coverageMap.cardinality();
                                //System.out.print(line);
                                //break;
                            } catch (NumberFormatException e) {
                                e.printStackTrace();
                            }
                        }
                        if (line.contains("AddressSanitizer")) {
                            coverage = 1;
                            resultStr = resultStr.concat(line + "\r\n");
                            System.out.print(resultStr);
                        }
                        resultStr = resultStr.concat(line + "\r\n");
                        //System.out.print(line);
                    }
                    if(resultStr.contains("Bitmap"))
                        break;
                }
            }
            catch(Exception e)
            {
                System.out.println(e);
            }
        }
        if(coverage!=1)
        {
            //System.out.print(coverage+" ");
            byte[] cov=new byte[8];
            fileCoverage.get(cov,0,8);
            fileCoverage.rewind();
            String tmp = new String(cov,"ASCII");
            coverage = Integer.valueOf(tmp.trim());
            //System.out.println(coverage);
        }
        return coverage;
    }
}

